import torch
import math
from typing import List
import torch.nn as nn
from torch import Tensor
import torch.nn.functional as F

BN_MOMENTUM=0.1


class Bottleneck(nn.Module):
    expansion = 0.5
    def __init__(self, inplanes, planes, stride=1, downsample=None, my_expansion=None):
        super(Bottleneck, self).__init__()
        # if my_expansion is not None:
        #     expansion = my_expansion
        # else:
        #     expansion = 2
        # logger.debug(my_expansion)
        # expansion = 0.5 if my_expansion is None else my_expansion
        self.conv1 = nn.Conv2d(inplanes, planes, kernel_size=1, bias=False)
        self.bn1 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
        self.conv2 = nn.Conv2d(planes, planes, kernel_size=3, stride=stride,
                               padding=1, bias=False)
        self.bn2 = nn.BatchNorm2d(planes, momentum=BN_MOMENTUM)
        self.conv3 = nn.Conv2d(planes, int(planes * self.expansion), kernel_size=1,
                               bias=False)
        self.bn3 = nn.BatchNorm2d(int(planes * self.expansion),
                                  momentum=BN_MOMENTUM)
        self.relu = nn.ReLU(inplace=True)
        self.downsample = downsample
        self.stride = stride

    def forward(self, x):
        residual = x

        out = self.conv1(x)
        out = self.bn1(out)
        out = self.relu(out)

        out = self.conv2(out)
        out = self.bn2(out)
        out = self.relu(out)

        out = self.conv3(out)
        out = self.bn3(out)

        if self.downsample is not None:
            residual = self.downsample(x)

        out += residual
        out = self.relu(out)

        return out


class TwoLayerConv2d(nn.Sequential):
    def __init__(self, in_channels, out_channels, kernel_size=3):
        super().__init__(nn.Conv2d(in_channels, in_channels, kernel_size=kernel_size,
                            padding=kernel_size // 2, stride=1, bias=False),
                         nn.BatchNorm2d(in_channels),
                         nn.ReLU(),
                         nn.Conv2d(in_channels, out_channels, kernel_size=kernel_size,
                            padding=kernel_size // 2, stride=1)
                         )


class OurPixelwiseLinear(nn.Module):
    def __init__(
        self,
        fin: List[int],
        fout: List[int],
        last_activation: nn.Module = None,
    ) -> None:
        assert len(fout) == len(fin)
        super().__init__()
        self.fout = fout
        self.fin = fin

        n = len(fin)
        self.n = n
        self._linears = nn.Sequential(
            *[
                nn.Sequential(
                    nn.Conv2d(fin[i], fout[i], kernel_size=1, bias=True),
                    nn.PReLU()
                    if i < n - 1 or last_activation is None
                    else last_activation,
                )
                for i in range(n-1)
            ]
        )
        # 卷积后有加BN，则bias=True
        self.final_conv = nn.Conv2d(fin[-1], fout[-1], kernel_size=1, bias=True)
        self.last_activation = last_activation
        self.mlp_layer = self._make_layer(n, fin, fout)
        

    def _make_layer(self, n, fin, fout):
        mlp_layer = []
        for i in range(n):
            if i == 0:
                mlp_layer.append(
                    nn.Sequential(
                        nn.Conv2d(
                            in_channels=fin[i],
                            out_channels=fout[i],
                            kernel_size=3, padding=1,
                            stride=1, bias=False
                        )
                    )
                )
            elif i <= n - 1:
                mlp_layer.append(
                    nn.Sequential(
                        nn.PReLU(),
                        nn.Conv2d(
                            in_channels=fin[i],
                            out_channels=fout[i],
                            kernel_size=3, padding=1,
                            stride=1, bias=False
                        )
                    )
                )
        return nn.ModuleList(mlp_layer)

    def forward(self, x: Tensor) -> Tensor:
        # Processing the tensor:
        # return self._linears(x)
        x_list = []
        for i in range(self.n):
            x = self.mlp_layer[i](x)
            x_list.append(x)
        if self.last_activation:
            x_list[-1] = self.last_activation(x_list[-1])
        return x_list


class CAU_MLP(nn.Module):
    def __init__(self, in_channels=32, out_channels=5, mlp_in_channels=[4, 4, 2], mlp_out_channels=[4, 2, 1], last_activation=nn.Sigmoid()):
        assert len(mlp_in_channels) == len(mlp_out_channels)
        super(CAU_MLP, self).__init__()
        self.inplanes = in_channels
        self.out_channels = out_channels
        self.classifier_1 = self._make_layer(Bottleneck, 16)
        self.classifier_2 = TwoLayerConv2d(in_channels=8, out_channels=self.out_channels)
        self.pseudo_gt = nn.Conv2d(in_channels=4, out_channels=2, kernel_size=1, padding=0, stride=1)
        self.multi_output_mlp = OurPixelwiseLinear(mlp_in_channels, mlp_out_channels, last_activation)


    def _make_layer(self, block, planes=16, blocks=1, stride=1):
        downsample = None
        if stride != 1 or self.inplanes != int(planes * block.expansion):
            downsample = nn.Sequential(
                nn.Conv2d(self.inplanes, int(planes * block.expansion),
                          kernel_size=1, stride=stride, bias=False),
                nn.BatchNorm2d(int(planes * block.expansion), momentum=BN_MOMENTUM),
            )

        layers = []
        layers.append(block(self.inplanes, planes, stride, downsample))
        self.inplanes = int(planes * block.expansion)
        for i in range(1, blocks):
            layers.append(block(self.inplanes, planes))

        return nn.Sequential(*layers)
    
    def _forward_pred_xor(self, x):
        x = F.softmax(x, dim=1)
        x = x.log()
        x = torch.argmax(x, dim=1)
        x = torch.unsqueeze(x, 1)
        return x
    
    def _forward_pseudo_gt(self, x, gt):
        
        t_gt = 2 * gt
        mask = t_gt + x
        mask = mask.type(torch.uint8)
        return mask
    
    def _forward_get_pseudo_gt(self, x, gt):
        x_xor = self.pseudo_gt(x)
        x_xor = self._forward_pred_xor(x_xor)
        pseudo_gt = self._forward_pseudo_gt(x_xor, gt)
        return pseudo_gt


    def forward(self, x, gt):
        x = self.classifier_1(x)
        x = self.classifier_2(x)
        offset_output = x[:, 4:, :, :].contiguous()
        x_output = x[:, :4, :, :].contiguous()
        pseudo_gt = None

        if gt is not None:
            pseudo_gt = self._forward_get_pseudo_gt(x_output, gt)

        mlp_3, mlp_2, mlp_1 = self.multi_output_mlp(x_output)
        final_output = offset_output + mlp_1
        return [final_output, offset_output, pseudo_gt, mlp_3, mlp_2, mlp_1]
        # return sum([[final_output, offset_output, pseudo_gt], x_list], [])